/**
 * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2018. ALL RIGHTS RESERVED.
 *
 * See file LICENSE for terms.
 */

#ifdef HAVE_CONFIG_H
#  include "config.h"
#endif

#include <tools/perf/lib/libperf_int.h>

#include "cnrt.h"
#include "cn_api.h" // CNresult
#include <ucs/sys/compiler.h>

static ucs_status_t ucx_perf_cnrt_init(ucx_perf_context_t *perf)
{
    cnrtRet_t cerr;
    unsigned group_index;
    int num_gpus;
    int gpu_index;

    group_index = rte_call(perf, group_index);

    cerr = cnrtGetDeviceCount(&num_gpus);
    if (cerr != cnrtSuccess) {
        return UCS_ERR_NO_DEVICE;
    }

    gpu_index = group_index % num_gpus;

    cerr = cnrtSetDevice(gpu_index);
    if (cerr != cnrtSuccess) {
        return UCS_ERR_NO_DEVICE;
    }

    /* actually set device context as calling cnrtSetDevice may result in
     * context being initialized lazily */
    cnrtFree(0);

    return UCS_OK;
}

static inline ucs_status_t ucx_perf_cnrt_alloc(size_t length,
                                               ucs_memory_type_t mem_type,
                                               uct_allocated_memory_t *alloc_mem)
{
    // CNresult ret;
    // ucs_warn("before alloc %p, %ld.", &(alloc_mem->mlu_addr), alloc_mem->mlu_addr);
    // ret = cnMallocPeerAble(&(alloc_mem->mlu_addr), length);
    // ucs_warn("after alloc %p, %ld.", &(alloc_mem->mlu_addr), alloc_mem->mlu_addr);
    // if (ret != CN_SUCCESS) {
    //     ucs_error("failed to allocate cncl device memory.");
    //     return UCS_ERR_NO_MEMORY;
    // }

    // alloc_mem->address = (void *)&(alloc_mem->mlu_addr);

    CNresult error;
    const size_t gpu_page_size = 64 * 1024;
    size_t size = (length + gpu_page_size - 1) &
        ~(gpu_page_size - 1);

    ucs_error("cnMallocPeerAble() of a %lu bytes MLU buffer\n",
            length);
    error = cnMallocPeerAble(&(alloc_mem->mlu_addr), size);
    if (error != CN_SUCCESS) {
        ucs_error("cnMallocPeerAble error=%d\n", error);
        return UCS_ERR_NO_MEMORY;
    }

    ucs_error("allocated %zu bytes of MLU buffer at %ld\n", size, alloc_mem->mlu_addr);
    alloc_mem->address = (void *)alloc_mem->mlu_addr;

    return UCS_OK;
}

static inline ucs_status_t
uct_perf_cnrt_alloc_reg_mem(const ucx_perf_context_t *perf,
                            size_t length,
                            ucs_memory_type_t mem_type,
                            unsigned flags,
                            uct_allocated_memory_t *alloc_mem)
{
    ucs_status_t status;


    status = ucx_perf_cnrt_alloc(length, mem_type, alloc_mem);

    if (status != UCS_OK) {
        ucs_error("failed to ucx_perf_cnrt_alloc memory");
        return status;
    }
    ucs_warn("test uct_perf_cnrt_alloc_reg_mem");
    status = uct_md_mem_reg(perf->uct.md, alloc_mem->address,
                            length, flags, &alloc_mem->memh);
    if (status != UCS_OK) {
        cnFree(alloc_mem->mlu_addr);
        alloc_mem->address = NULL;
        ucs_error("failed to register memory");
        return status;
    }

    alloc_mem->mem_type = mem_type;
    alloc_mem->md       = perf->uct.md;

    return UCS_OK;
}

static ucs_status_t uct_perf_cnrt_alloc(const ucx_perf_context_t *perf,
                                        size_t length, unsigned flags,
                                        uct_allocated_memory_t *alloc_mem)
{
    return uct_perf_cnrt_alloc_reg_mem(perf, length, UCS_MEMORY_TYPE_CNCL,
                                       flags, alloc_mem);
}

static void uct_perf_cnrt_free(const ucx_perf_context_t *perf,
                               uct_allocated_memory_t *alloc_mem)
{
    ucs_status_t status;
    CNaddr mlu_addr = (CNaddr)alloc_mem->mlu_addr;

    ucs_assert(alloc_mem->md == perf->uct.md);

    status = uct_md_mem_dereg(perf->uct.md, alloc_mem->memh);
    if (status != UCS_OK) {
        ucs_error("failed to deregister memory");
    }

    printf("deallocating RX MLU buffer %ld\n", mlu_addr);
    cnFree(mlu_addr);
    // cnFree(alloc_mem->mlu_addr);
    alloc_mem->address = NULL;
}

static void ucx_perf_cnrt_memcpy(void *dst, ucs_memory_type_t dst_mem_type,
                                 const void *src, ucs_memory_type_t src_mem_type,
                                 size_t count)
{
    CNaddr dst_a = (CNaddr)dst;
    CNaddr src_a = (CNaddr)src;
    CNresult ret = cnMemcpy(dst_a, src_a, (cn_uint64_t)count+1); // 结束符
    
    if (ret != CN_SUCCESS) {
        ucs_error("failed to cp memory %d", ret);
        return;
    }
}

static void* ucx_perf_cnrt_memset(void *dst, int value, size_t count)
{
    CNaddr dst_a = (CNaddr)dst;
    CNresult ret = cnMemsetD8(dst_a, value, count);
    if (ret != CN_SUCCESS) {
        const char** str = NULL;
        cnGetErrorString(ret, str);
        ucs_error("failed to set memory: %s", *str);
    }

    return dst;
}

UCS_STATIC_INIT {
    static ucx_perf_allocator_t cnrt_allocator = {
        .mem_type  = UCS_MEMORY_TYPE_CNCL,
        .init      = ucx_perf_cnrt_init,
        .uct_alloc = uct_perf_cnrt_alloc,
        .uct_free  = uct_perf_cnrt_free,
        .memcpy    = ucx_perf_cnrt_memcpy,
        .memset    = ucx_perf_cnrt_memset
    };

    ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CNCL]         = &cnrt_allocator;
}
UCS_STATIC_CLEANUP {
    ucx_perf_mem_type_allocators[UCS_MEMORY_TYPE_CNCL]         = NULL;
}
